import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
from sklearn import svm
import joblib  # 用于保存 SVM 模型
import numpy as np

class MNISTDataset(Dataset):
    def __init__(self, img_dir, labels_file, transform=None):
        self.img_dir = img_dir
        self.labels_df = pd.read_excel(labels_file)
        self.transform = transform

    def __len__(self):
        return len(self.labels_df)

    def __getitem__(self, idx):
        img_name = self.labels_df.iloc[idx, 0] + '.jpg'
        label = self.labels_df.iloc[idx, 1]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path).convert('L')  # Convert image to grayscale
        if self.transform:
            image = self.transform(image)
        return image, label

# Define your transformation: resize and normalization
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # MNIST normalization values
])

dataset = MNISTDataset(img_dir='./data/train_data', labels_file='./data/train_data_label.xlsx', transform=transform)
data_loader = DataLoader(dataset, batch_size=64, shuffle=True)

# 将数据转换为适合 SVM 输入的格式
X_train = []
y_train = []

for images, labels in data_loader:
    images = images.view(images.size(0), -1)  # 将每张图片展平成一维
    X_train.extend(images.numpy())  # 将 torch tensor 转换为 numpy array
    y_train.extend(labels.numpy())

X_train = np.array(X_train)
y_train = np.array(y_train)

# 初始化 SVM
model = svm.SVC(kernel='linear')

# 训练 SVM
print("Training SVM... ")
model.fit(X_train, y_train)
print("Training complete.")

# 保存 SVM 模型
joblib.dump(model, 'mnist_svm.pkl')
print('Model saved to mnist_svm.pkl')

# 通过交叉验证来评估模型性能
from sklearn.model_selection import cross_val_score
scores = cross_val_score(model, X_train, y_train, cv=5)  # 5-fold cross-validation
print(f'Cross-validation accuracy: {scores.mean() * 100:.3f}%')
